'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 2_study_loop_and_condition.py
* @Time: 2025/11/4
* @All Rights Reserve By Brtc
'''
import json
from typing import TypedDict, Annotated, Any, Literal

import dotenv
from langchain_community.tools import GoogleSerperRun
from langchain_community.tools.openai_dalle_image_generation import OpenAIDALLEImageGenerationTool
from langchain_community.utilities import GoogleSerperAPIWrapper
from langchain_community.utilities.dalle_image_generator import DallEAPIWrapper
from langchain_core.messages import ToolMessage
from langchain_openai import ChatOpenAI
from langgraph.constants import END, START
from langgraph.graph import add_messages, StateGraph
from pydantic import BaseModel, Field

dotenv.load_dotenv()
class GoogleSerperSchema(BaseModel):
    query:str = Field(description="执行谷歌搜索的查询语句")

google_serper = GoogleSerperRun(
    name = "google_serper",
    description = ("一个低成本的谷歌搜索API工具！"),
    args_schema=GoogleSerperSchema,
    api_wrapper=GoogleSerperAPIWrapper(),)

dalle = OpenAIDALLEImageGenerationTool(
    name="openai-dalle",
    api_wrapper=DallEAPIWrapper(model = "dall-e-3")
)
# 1、定义状态数据
class ChatBotState(TypedDict):
    messages:Annotated[list, add_messages]

bot_tools = [google_serper, dalle]
llm = ChatOpenAI(model="gpt-4o-mini")
llm_with_tools = llm.bind_tools(bot_tools)

def chat_bot(state:ChatBotState)->ChatBotState:
    """聊天机器人"""
    ai_meaasge = llm_with_tools.invoke(state["messages"])
    return {"messages":[ai_meaasge]}

def tool_exe(state:ChatBotState)->Any:
    tool_by_name = {tool.name:tool for tool in bot_tools }
    tool_calls = state["messages"][-1].tool_calls
    messages = []
    for tool_call in tool_calls:
        tool = tool_by_name[tool_call["name"]]
        messages.append(ToolMessage(
                tool_call_id=tool_call["id"],
                content=json.dumps(tool.invoke(tool_call["args"])),
                name = tool_call["name"],
            ))
    return {"messages":messages}

def route(state:ChatBotState)->Literal["tool_exe", "__end__"]:
    """动态选择工具亦或者直接结束"""
    ai_message = state["messages"][-1]
    if hasattr(ai_message, "tool_calls") and len(ai_message.tool_calls) > 0:
        return "tool_exe"
    return END

"""    
图的架构
             ----------tool_exe   
            |             |
            |             |
开始------>chat_bot----->route------->结束

"""



#1、定义图的状态结构
graph_builder = StateGraph(ChatBotState)
#2、添加节点
graph_builder.add_node("chat_bot", chat_bot)
graph_builder.add_node("tool_exe", tool_exe)
#3、添加边
graph_builder.add_edge(START, "chat_bot")
graph_builder.add_conditional_edges("chat_bot", route)
graph_builder.add_edge("tool_exe", "chat_bot")
#4、编译图为可运行组件
grap_app = graph_builder.compile()

state = grap_app.invoke({"messages":[("human", "2024年北京半程马拉松的前三名是谁？分别是什么国籍？")]})
for message in state["messages"]:
    print(f"消息类型:{message.type}")
    if hasattr(message, "tool_calls") and len(message.tool_calls) > 0:
        print(f"调用工具:{message.tool_calls}")
    print(f"消息内容:{message.content}")
    print("==================================================")



